from typing import AsyncIterator, Iterator, Any

from langchain_community.document_loaders import FileSystemBlobLoader
from langchain_community.document_loaders.generic import GenericLoader
from langchain_core.document_loaders import BaseLoader, BaseBlobParser, Blob
from langchain_core.documents import Document
import asyncio


# 基于LLM的应用程序通常涉及从数据库或文件(如PDF)中提取数据，并将其转换为LLM可以利用的格式。
# 在LangChain中，这通常涉及创建Document对象，该对象封装了提取的文本(page_content)以及元数据一包含有关文档的详细信息的字典，例如作者姓名或出版日期。
# Document 对象通常被格式化为提示，然后输入LLM以便LLM可以使用 Document 中的信息生成所需的响应(例如，对文档进行摘要)。
# Documents 可以立即使用，也可以索引到向量存储中以供将来检索和使用。文档加载的主要抽象为:
#
#     Document : 包含 text 和 metadata 的内容
#     BaseLoader : 文档加载器 ，描述如何从原始数据源加载数据
#     Blob : 存储原始数据的二进制数据
#     BaseBlobParser : 二进制数据解析器，描述如何将 Blob 转换为 Document

# 自定义文档加载器
class CustomDocumentLoader(BaseLoader):
    """一个从文件逐行读取的示例文档加载器。"""

    def __init__(self, file_path: str) -> None:
        """使用文件路径初始化加载器。
        Args:
            file_path: 要加载的文件的路径。
        """
        self.file_path = file_path

    def lazy_load(self) -> Iterator[Document]:  # <-- 不接受任何参数
        """逐行读取文件的惰性加载器。
        当您实现惰性加载方法时，应使用生成器逐个生成文档。
        """
        with open(self.file_path, encoding="utf-8") as f:
            line_number = 0
            for line in f:
                # yield 关键字用于定义生成器函数。
                # 生成器函数是一种特殊类型的函数，它允许你逐步生成一个序列的值，而不是一次性返回整个序列。
                # 与普通的函数不同，生成器函数在每次调用时不会从头开始执行，而是从上次离开的地方继续执行。
                # 这使得生成器非常适合处理需要逐步生成或处理大数据集的情况
                yield Document(
                    page_content=line,
                    metadata={"line_number": line_number, "source": self.file_path},
                )
                line_number += 1

    # alazy_load是可选的。
    # 如果您省略了实现，将使用默认实现，该实现将委托给lazy_load！
    async def alazy_load(self) -> AsyncIterator[Document]:  # <-- 不接受任何参数
        """逐行读取文件的异步惰性加载器。"""
        # 需要aiofiles
        # 使用`pip install aiofiles`安装
        # https://github.com/Tinche/aiofiles
        import aiofiles
        async with aiofiles.open(self.file_path, encoding="utf-8") as f:
            line_number = 0
            async for line in f:
                # async for 逐行异步读取文件内容，并通过 yield 关键字逐行生成 Document 对象。
                # 异步生成器允许你在异步环境中逐步生成值，这对于处理 I/O 密集型任务（如文件读取或网络请求）非常有用，
                # 因为它可以在等待 I/O 操作完成时释放事件循环，以便处理其他任务。
                yield Document(
                    page_content=line,
                    metadata={"line_number": line_number, "source": self.file_path},
                )
                line_number += 1


loader = CustomDocumentLoader("./meow.txt")

# # 测试延迟加载接口
# for doc in loader.lazy_load():
#     print('==========================')
#     print(type(doc))
#     print(doc)


# async def alazy():
#     async for doc in loader.alazy_load():
#         print()
#         print(type(doc))
#         print(doc)
#
# # 测试异步实现
# asyncio.run(alazy())

# load() 在诸如 Jupyter Notebook 之类的交互式环境中很有用。
# 在生产代码中避免使用它，因为急切加载假定所有内容都可以放入内存中，
# 而这并不总是成立，特别是对于企业数据而言
# print(loader.load())
print("==============================================================")


# 自定义文档解析器
class MyParser(BaseBlobParser):
    """一个简单的解析器，每行创建一个文档。"""

    def lazy_parse(self, blob: Blob) -> Iterator[Document]:
        """逐行将 blob 解析为文档。"""
        line_number = 0
        with blob.as_bytes_io() as f:
            for line in f:
                line_number += 1
                yield Document(
                    page_content=line,
                    metadata={"line_number": line_number, "source": blob.source},
                )


# blob = Blob.from_path("./meow.txt", encoding="utf-8")
# parser = MyParser()
# print(list(parser.lazy_parse(blob)))

# blob 允许直接从内存加载内容
# blob = Blob(data="来自内存的一些数据\n喵".encode("utf-8"), encoding="utf-8")
# parser = MyParser()
# print(list(parser.lazy_parse(blob)))
print("==============================================================")

# 2.blob API
blob = Blob.from_path("./meow.txt", metadata={"foo": "bar"})

# print('编码 ： ', blob.encoding)
# print('字节数据 ： ', blob.as_bytes())
# print('字符串数据 ： ', blob.as_string())
# print('io对象 ： ', blob.as_bytes_io())
#
# # Blob 元数据
# print('元信息 ： ', blob.metadata)
# print('来源 ： ', blob.source)

print("==============================================================")


# 自定义 blob 加载器
class BlobParser(BaseBlobParser):
    """一个简单的解析器，每行创建一个文档。"""

    def lazy_parse(self, blob: Blob) -> Iterator[Document]:
        """逐行将 blob 解析为文档。"""
        line_number = 0
        with blob.as_bytes_io() as f:
            for line in f:
                line_number += 1
                yield Document(
                    page_content=line,
                    metadata={"line_number": line_number, "source": blob.source},
                )


# parser2 = BlobParser()
# blob_loader = FileSystemBlobLoader(path="./resource", glob="*.mdx", show_progress=True)
# for blob in blob_loader.yield_blobs():
#     for doc in parser2.lazy_parse(blob):
#         print(doc)
#         break
print("==============================================================")

# 通用加载器
generic_loader = GenericLoader.from_filesystem(path="./resource", glob="*.mdx", show_progress=True,
                                               parser=MyParser())

# for idx, doc in enumerate(generic_loader.lazy_load()):
#     print(doc)
# print("... output truncated for demo purposes")

print("==============================================================")


# 自定义
class MyCustomLoader(GenericLoader):
    @staticmethod
    def get_parser(**kwargs: Any) -> BaseBlobParser:
        """Override this method to associate a default parser with the class."""
        return MyParser()


# loader3 = MyCustomLoader.from_filesystem(path="./resource", glob="*.mdx", show_progress=True)
# for idx, doc in enumerate(loader3.lazy_load()):
#     print(doc)
